// Created by Chengfu Zou
// Copyright (C) FYT Vision Group. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "rm_utils/math/trajectory_compensator.hpp"
#include<iostream>

namespace fyt {
bool TrajectoryCompensator::compensate(const Eigen::Vector3d &target_position,
                                       double &pitch) const noexcept {
    double target_height = target_position.z();
    double distance = std::hypot(target_position.x(), target_position.y());

    // 用 Newton 法直接求解仰角
    auto f  = [&](double ang){
        return calculateTrajectory(distance, ang) - target_height;
    };
    auto df = [&](double ang){
        const double d = 1e-5;
        return (f(ang + d) - f(ang - d)) / (2*d);
    };

    double angle = std::atan2(target_height, distance);
    double delta = 1.0;
    for (int i = 0; i < iteration_times && std::abs(delta) > 1e-3; ++i) {
        double yv = f(angle), yvp = df(angle);
        if (std::abs(yvp) < 1e-8) break;
        delta = yv / yvp;
        angle -= delta;
    }

    if (std::abs(f(angle)) > 0.01 || std::abs(angle) > M_PI/2.5)
        return false;
    pitch = angle;
    return true;}

std::vector<std::pair<double, double>> TrajectoryCompensator::getTrajectory(
  double distance, double angle) const noexcept {
  std::vector<std::pair<double, double>> trajectory;

  if (distance < 0) {
    return trajectory;
  }

  for (double x = 0; x < distance; x += 0.03) {
    trajectory.emplace_back(x, calculateTrajectory(x, angle));
  }
  return trajectory;
}

double IdealCompensator::calculateTrajectory(const double x, const double angle) const noexcept {
  double t = x / (velocity * cos(angle));
  double y = velocity * sin(angle) * t - 0.5 * gravity * t * t;
  return y;
}

double IdealCompensator::getFlyingTime(const Eigen::Vector3d &target_position) const noexcept {
  double distance =
    sqrt(target_position(0) * target_position(0) + target_position(1) * target_position(1));
  double angle = atan2(target_position(2), distance);
  double t = distance / (velocity * cos(angle));
  return t;
}


double ResistanceCompensator::calculateTrajectory(const double x,
                                                  const double angle) const noexcept {
    struct State { double x, y, vx, vy; };
    auto rk4Step = [&](const State &s, double h) {
        auto deriv = [&](const State &u) {
            double v = std::hypot(u.vx, u.vy);
            // 动态气密度 & Cd
            double rho = airDensity(u.y);
            double Cd  = dragCoefficient(v);
            // —— 修改：真正带上 Cd 的阻力加速度计算 —— 
            double a_drag = 0.5 * rho * Cd * projectileArea * v * v / projectileMass;
            // Magnus 力（横向）
            double a_magnus = (spinAngularRate != 0.0)
                ? (0.5 * rho * v * spinAngularRate * projectileArea / projectileMass)
                : 0.0;

            double ax = -a_drag * (u.vx / v)
                        +  a_magnus * (u.vy / v);
            double ay = -gravity
                        - a_drag * (u.vy / v)
                        -  a_magnus * (u.vx / v);
            return State{u.vx, u.vy, ax, ay};
        };
        State k1 = deriv(s);
        State k2 = deriv(State{s.x + 0.5*h*k1.x,
                               s.y + 0.5*h*k1.y,
                               s.vx + 0.5*h*k1.vx,
                               s.vy + 0.5*h*k1.vy});
        State k3 = deriv(State{s.x + 0.5*h*k2.x,
                               s.y + 0.5*h*k2.y,
                               s.vx + 0.5*h*k2.vx,
                               s.vy + 0.5*h*k2.vy});
        State k4 = deriv(State{s.x +   h*k3.x,
                               s.y +   h*k3.y,
                               s.vx +   h*k3.vx,
                               s.vy +   h*k3.vy});
        return State{
            s.x  + (k1.x  + 2*k2.x  + 2*k3.x  + k4.x ) * (h/6.0),
            s.y  + (k1.y  + 2*k2.y  + 2*k3.y  + k4.y ) * (h/6.0),
            s.vx + (k1.vx + 2*k2.vx + 2*k3.vx + k4.vx) * (h/6.0),
            s.vy + (k1.vy + 2*k2.vy + 2*k3.vy + k4.vy) * (h/6.0)
        };
    };

    State s{0.0, 0.0,
            velocity * std::cos(angle),
            velocity * std::sin(angle)};
    const double h = 1e-3;
    State prev = s;

    // 主循环：RK4 积分
    while (s.x < x && s.y >= 0.0) {
        prev = s;
        s = rk4Step(s, h);
    }

    // Hermite 插值：提高越界后的高度精度
    double dx = s.x - prev.x;
    if (dx <= 0.0) return prev.y;
    double t = (x - prev.x) / dx;

    // Hermite 基函数
    double y0 = prev.y, y1 = s.y;
    double m0 = (prev.vy / prev.vx) * dx;
    double m1 = (s.vy   / s.vx)   * dx;
    double t2 = t*t, t3 = t2*t;
    double h00 =  2*t3 - 3*t2 + 1;
    double h10 =    t3 - 2*t2 + t;
    double h01 = -2*t3 + 3*t2;
    double h11 =    t3 -   t2;

    return h00*y0 + h10*m0 + h01*y1 + h11*m1;
}


double ResistanceCompensator::getFlyingTime(const Eigen::Vector3d &target_position) const noexcept {
    // 水平距离 & 目标高度
    const double x_target = std::hypot(target_position.x(), target_position.y());
    const double y_target = target_position.z();

    // 使用 Newton 法求解 f(angle) = calculateTrajectory(x_target, angle) - y_target = 0
    auto f  = [&](double ang){
        return calculateTrajectory(x_target, ang) - y_target;
    };
    auto df = [&](double ang){
        const double d = 1e-5;
        return (f(ang + d) - f(ang - d)) / (2*d);
    };

    // 初值：简单 atan2
    double angle = std::atan2(y_target, x_target);
    for (int i = 0; i < 20; ++i) {
        double yv = f(angle), yvp = df(angle);
        if (std::abs(yvp) < 1e-8) break;
        double delta = yv / yvp;
        angle -= delta;
        if (std::abs(delta) < 1e-7) break;
    }

    // 最终积分累加时间（RK4，不含 Hermite）
    struct State { double x, y, vx, vy; };
    auto rk4Step = [&](const State &s, double h) {
        // 同 calculateTrajectory 中的 deriv
        auto deriv = [&](const State &u) {
            double v = std::hypot(u.vx, u.vy);
            double rho = airDensity(u.y);
            double Cd  = dragCoefficient(v);
            // —— 修改：真正带上 Cd 的阻力加速度计算 —— 
            double a_drag = 0.5 * rho * Cd * projectileArea * v * v / projectileMass;
            double a_magnus = (spinAngularRate != 0.0)
                ? (0.5 * rho * v * spinAngularRate * projectileArea / projectileMass)
                : 0.0;
            double ax = -a_drag * (u.vx / v)
                        +  a_magnus * (u.vy / v);
            double ay = -gravity
                        - a_drag * (u.vy / v)
                        -  a_magnus * (u.vx / v);
            return State{u.vx, u.vy, ax, ay};
        };
        State k1 = deriv(s);
        State k2 = deriv(State{s.x + 0.5*h*k1.x,
                               s.y + 0.5*h*k1.y,
                               s.vx + 0.5*h*k1.vx,
                               s.vy + 0.5*h*k1.vy});
        State k3 = deriv(State{s.x + 0.5*h*k2.x,
                               s.y + 0.5*h*k2.y,
                               s.vx + 0.5*h*k2.vx,
                               s.vy + 0.5*h*k2.vy});
        State k4 = deriv(State{s.x +   h*k3.x,
                               s.y +   h*k3.y,
                               s.vx +   h*k3.vx,
                               s.vy +   h*k3.vy});
        return State{
            s.x  + (k1.x  + 2*k2.x  + 2*k3.x  + k4.x ) * (h/6.0),
            s.y  + (k1.y  + 2*k2.y  + 2*k3.y  + k4.y ) * (h/6.0),
            s.vx + (k1.vx + 2*k2.vx + 2*k3.vx + k4.vx) * (h/6.0),
            s.vy + (k1.vy + 2*k2.vy + 2*k3.vy + k4.vy) * (h/6.0)
        };
    };

    State s{0.0, 0.0,
            velocity * std::cos(angle),
            velocity * std::sin(angle)};
    const double h = 1e-3;
    double time = 0.0;
    while (s.x < x_target && s.y >= 0.0) {
        s = rk4Step(s, h);
        time += h;
    }
    return time;
}



}  // namespace fyt
